#!/bin/bash

stage=5

if [ $stage -eq 0 ]; then
    export PATH="/ceph/home/zhangy20/dpss-exp3-VC-PPG/SPTK-3.11/build/bin:$PATH"

    CUDA_VISIBLE_DEVICES=0 TF_FORCE_GPU_ALLOW_GROWTH=true \
        python preprocess.py \
        --data_dir wav/cmu_us_slt_arctic \
        --save_dir save/cmu-arctic-slt/ &

    CUDA_VISIBLE_DEVICES=1 TF_FORCE_GPU_ALLOW_GROWTH=true \
        python preprocess.py \
        --data_dir wav \
        --save_dir save/cmu_arctic-all/
    wait
fi

# Any2One Train
if [ $stage -eq 1 ]; then
    echo Any2One Train

    # ResNet Training
    rm -rf res.log
    CUDA_VISIBLE_DEVICES=1 TF_FORCE_GPU_ALLOW_GROWTH=true \
        python train_to_one.py \
        --model_dir ./ckpt/Any2One/res \
        --test_dir ./test/Any2One/res \
        --data_dir ./save/cmu-arctic-slt/ \
        --ResNet 

    # NoResNet Training
    #rm -rf nores.log
    #CUDA_VISIBLE_DEVICES=2 TF_FORCE_GPU_ALLOW_GROWTH=true \
    #    python train_to_one.py \
    #    --model_dir ./ckpt/Any2One/nores\
    #    --test_dir ./test/Any2One/nores \
    #    --data_dir ./save/cmu-arctic-slt/ 
fi


# Any2One Inference
if [ $stage -eq 2 ]; then
    rm -rf test_dir
    mkdir test_dir

    wav_path=test.wav

    # No Res
    CUDA_VISIBLE_DEVICES=0 TF_FORCE_GPU_ALLOW_GROWTH=true \
        python inference_to_one.py --src_wav $wav_path \
        --ckpt ckpt/Any2One/nores/ppg-vc-to-one-59.pt \
        --save_dir ./test_dir/

    # ResNet
    CUDA_VISIBLE_DEVICES=0 TF_FORCE_GPU_ALLOW_GROWTH=true \
        python inference_to_one.py --src_wav $wav_path \
        --ckpt ckpt/Any2One/res/ppg-vc-to-one-59.pt \
        --save_dir ./test_dir/ \
        --ResNet

    cd test_dir
    ln -s ../test.wav .
    python -m http.server
fi


if [ $stage -eq 3 ]; then
    CUDA_VISIBLE_DEVICES=3 TF_FORCE_GPU_ALLOW_GROWTH=true \
        python train_to_many.py \
        --model_dir ./ckpt/Any2Many \
        --test_dir ./test/Any2Many \
        --data_dir ./save/cmu_arctic-all/
fi


if [ $stage -eq 4 ]; then
    rm -rf test_dir
    mkdir test_dir

    for tgt in awb bdl clb jmk rms slt; do
        echo "################################"
        echo $tgt
        echo "################################"
        CUDA_VISIBLE_DEVICES=3 TF_FORCE_GPU_ALLOW_GROWTH=true \
            python inference_to_many.py \
            --src_wav test.wav \
            --tgt_spk $tgt \
            --ckpt ckpt/Any2Many/ppg-vc-to-many-149.pt \
            --save_dir ./test_dir/
    done

    cd test_dir
    #ln -s ../../speaker_embedding/dataset/* .
    ln -s ../test.wav .
    python -m http.server
fi


if [ $stage -eq 5 ]; then
    rm -rf test_dir
    mkdir test_dir

    CUDA_VISIBLE_DEVICES=3 TF_FORCE_GPU_ALLOW_GROWTH=true \
        python inference_to_many.py \
        --src_wav input1.wav \
        --tgt_spk slt \
        --ckpt ckpt/Any2Many/ppg-vc-to-many-149.pt \
        --save_dir ./test_dir/

    CUDA_VISIBLE_DEVICES=3 TF_FORCE_GPU_ALLOW_GROWTH=true \
        python inference_to_many.py \
        --src_wav input2.wav \
        --tgt_spk rms \
        --ckpt ckpt/Any2Many/ppg-vc-to-many-149.pt \
        --save_dir ./test_dir/

    CUDA_VISIBLE_DEVICES=3 TF_FORCE_GPU_ALLOW_GROWTH=true \
        python inference_to_many.py \
        --src_wav input3.wav \
        --tgt_spk awb \
        --ckpt ckpt/Any2Many/ppg-vc-to-many-149.pt \
        --save_dir ./test_dir/

    cd test_dir
    cp ../input*.wav .
    python -m http.server
fi

